library(data.table)
library(gridExtra)
library(grid)
library(dplyr)
library(tidyr)
library(ggplot2)
library(stringr)
library(readr)
library(latex2exp)

## Paper or presentation
paper <- TRUE

## data directory
dir.project <- './'
dir.data <- paste0(dir.project, 'data/')
dir.raw <- paste0(dir.data, 'raw/')
dir.generated <- paste0(dir.data, 'generated/')

## figures directory for output
dir.fig.paper <- paste0(dir.project, 'doc/paper/figures/GPOPS/')
dir.fig.presentation <-
    paste0(dir.project, 'results/figures/presentation/')

source(paste0(dir.project, 'src/r/figures.R'))

fileName <- paste0(dir.project, 'calibration_version.txt')
calibration_version <-
    read_file(fileName) %>% str_replace_all("[\r\n]" , "")

#############
## FIGURES ##
#############

dat.optimum <-
    fread(paste0(dir.generated, 'matlab/OptimTax/gpops/', 'optimum_vars.csv')) %>%
    mutate(pricedrop = drop * 100)

dat.optimum_vectors <-
    fread(paste0(
        dir.generated,
        'matlab/OptimTax/gpops/',
        'optimum_vectors.csv'
    )) %>%
    mutate(pricedrop = drop * 100)



dat.optimum_long <- dat.optimum %>%
    select(pricedrop, tau_B, tau_E, tau_S) %>%
    mutate(pricedrop = as.numeric(pricedrop)) %>%
    gather(starts_with('tau_'), key='tax', value='value')

gpops_capital_taxes <- ggplot(dat.optimum_long , aes(x = pricedrop,y=value*100)) +
    geom_line(aes(linetype=tax)) +
    theme_bw() +
    theme.serif +
    theme(legend.position = c(0.2, 0.8),
          legend.background = element_rect(fill = alpha('white', 0))) +
        xlab('pct. drop in robot price') +
    ylab('') +
    scale_y_continuous(sec.axis = sec_axis( ~ . * 1.0436/0.0436, name = ''), limits = c(-1, 4)) +
    scale_linetype_manual(
        name = '',
        values = c('tau_B' = 1, 'tau_E' = 2, 'tau_S' = 3),
        labels = unname(TeX(c(
            "$t_B$", "$t_E$", "$t_S$"
        )))
    ) 


to.pdf(
    print(gpops_capital_taxes),
    "gpops_capital_taxes.pdf",
    width = width.double,
    height = height.double
)

dat.optimum <-
    dat.optimum %>% mutate(labor_income = (Y_M * L_M + Y_R * L_R + Y_C * L_C) /
                               part)

gpops_transfer <-
    ggplot(dat.optimum, aes(x = pricedrop, y = b / labor_income * 100)) + geom_line() +
    xlab('pct. drop in robot price') +
    ylab('') +
    theme_bw() +
    theme.serif

to.pdf(print(gpops_transfer),
       "gpops_transfer.pdf",
       width = width.double,
       height = height.double)

plots_combined <-
    grid.arrange(
        gpops_capital_taxes + ggtitle('A: Tax on capital in pct.'),
        gpops_transfer + ggtitle('B: Transfer (pct. of avg. lab. inc.)'),
        ncol = 2,
        nrow = 1,
        widths=c(7.3,6.2)
    )

ggsave(
    paste0(dir.fig, "gpops_plots_combined.pdf"),
    plot = plots_combined,
    width = 13.5,
    height = 6,
    units = "cm"
)


pd.ub <- max(dat.optimum$pricedrop)

# plot of marginal tax rates - still need to discuss how to deal with them
gpops_mtax_compare <- ggplot(
    dat.optimum_vectors %>%
        filter(pricedrop %in% c(0, 50, pd.ub), y <= 1000, mtax >
                   0),
    aes(x = y, y = mtax * 100)
) +
    geom_line(aes(linetype = as.factor(pricedrop), color = as.factor(pricedrop))) +
    scale_linetype_manual(
        name = 'pct. drop in robot price',
        values = c(1, 2, 3),
        labels = c(0, 50, pd.ub)
    ) +
    scale_color_manual(name = 'pct. drop in robot price',
                       values = myPalette,
                       labels = c(0, 50, pd.ub)) +
    xlab('labor inc. in 1000 USD') +
    ylab('') +
    theme_bw() +
    theme.serif +
    theme(legend.position = 'none')

to.pdf(
    print(gpops_mtax_compare),
    "gpops_mtax_compare.pdf",
    width = width.double,
    height = height.double
)

# plot of avg tax rates - still need to discuss how to deal with them
p <- ggplot(dat.optimum_vectors %>%
                filter(pricedrop %in% c(0, 50, pd.ub),y<=1000),aes(x = y, y = avg_tax * 100)) +
    geom_line(aes(linetype = as.factor(pricedrop), color = as.factor(pricedrop))) +
    scale_linetype_manual(
        name = 'pct. drop in robot price',
        values = c(1, 2, 3),
        labels = c(0, 50, pd.ub)) +
    scale_color_manual(name = 'pct. drop in robot price',
                       values = myPalette,
                       labels = c(0, 50, pd.ub)) +
    xlab('labor inc. in 1000 USD') +
    ylab('') +
    ylim(-400, 100) +
    theme_bw() +
    theme.serif +
    theme(legend.position = 'top')

leg <- p
gpops_avgtax_compare <- p + theme(legend.position = 'none')

to.pdf(print(p),
       "gpops_avgtax_compare.pdf",
       width = width.double,
       height = height.double)

# legend
to.pdf(fig.legend(leg),
       "legend_tax_compare.pdf",
       width = 5,
       height = 0.3)

plots_combined <-
    grid.arrange(
        gpops_mtax_compare + ggtitle('A: Marg. inc. tax rate in pct.'),
        gpops_avgtax_compare + ggtitle('B: Avg. inc. tax rate in pct.'),
        ncol = 2,
        nrow = 1
    )

ggsave(
    paste0(dir.fig, "gpops_tax_compare_combined.pdf"),
    plot = plots_combined,
    width = 13,
    height = 6,
    units = "cm"
)
